package org.hairball.graph;

import org.hairball.ClassInclusionFilter;

import java.util.HashSet;
import java.util.Set;

public class NodeFilteringGraph implements HairballGraph {
    private HairballGraph underlyingGraph;
    private ClassInclusionFilter filter;

    public NodeFilteringGraph(HairballGraph underlyingGraph, ClassInclusionFilter filter) {
        this.underlyingGraph = underlyingGraph;
        this.filter = filter;
    }

    public void visitEdges(final EdgeVisitor edgeVisitor) {
        Set<Edge> filteredEdges = filteredEdges();
        for (Edge filteredEdge : filteredEdges) {
            edgeVisitor.visit(filteredEdge.fromNode, filteredEdge.toNode, filteredEdge.type);
        }
    }

    public void visitNodes(final NodeVisitor nodeVisitor) {
        underlyingGraph.visitNodes(new NodeVisitor() {

            public void visit(String node, boolean connected) {
                if (filter.shouldInclude(node) && isConnected(node)) {
                    nodeVisitor.visit(node, connected);
                }
            }
        });
    }

    private Set<Edge> filteredEdges() {
        final Set<Edge> filteredEdges = new HashSet<Edge>();

        underlyingGraph.visitEdges(new EdgeVisitor() {

            public void visit(String fromNode, String toNode, String edgeType) {
                if (filter.shouldInclude(fromNode) && filter.shouldInclude(toNode)) {
                    filteredEdges.add(new Edge(fromNode, toNode, edgeType));
                }
            }
        });
        return filteredEdges;
    }

    private boolean isConnected(String node) {
        Set<Edge> edges = filteredEdges();

        for (Edge edge : edges) {
            if (edge.fromNode.equals(node) || edge.toNode.equals(node)) {
                return true;
            }
        }
        return false;
    }

}
